{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "JndnmDMp66FL"
      },
      "source": [
        "##### Copyright 2020 Google LLC.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "both",
        "colab": {},
        "colab_type": "code",
        "id": "hMqWDc_m6rUC"
      },
      "outputs": [],
      "source": [
        "#@title Default title text\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "pW8P5zUZFjPQ"
      },
      "source": [
        "# Train your own Keyword Spotting Model.\n",
        "[Open in Google  Colab](https://colab.research.google.com/github/google-research/google-research/blob/master/speech_embedding/record_train.ipynb)\n",
        "\n",
        "Before running any cells please enable GPUs for this notebook to speed it up. \n",
        "\n",
        "* *Edit* → *Notebook Settings*\n",
        "* select *GPU* from the *Hardware Accelerator* drop-down\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "X4vndgdGztBS"
      },
      "outputs": [],
      "source": [
        "#@title Imports\n",
        "%tensorflow_version 1.x\n",
        "from __future__ import division\n",
        "\n",
        "import collections\n",
        "import IPython\n",
        "import functools\n",
        "import math\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import io\n",
        "import os\n",
        "import tensorflow as tf\n",
        "import tensorflow_hub as hub\n",
        "import random\n",
        "import scipy.io.wavfile\n",
        "import tarfile\n",
        "import time\n",
        "import sys\n",
        "\n",
        "from google.colab import output\n",
        "from google.colab import widgets\n",
        "from base64 import b64decode\n",
        "\n",
        "!pip install ffmpeg-python\n",
        "import ffmpeg"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "3C3TL3WbztBd"
      },
      "outputs": [],
      "source": [
        "#@title Helper functions and classes\n",
        "def normalized_read(filename):\n",
        "  \"\"\"Reads and normalizes a wavfile.\"\"\"\n",
        "  _, data = scipy.io.wavfile.read(open(filename, mode='rb'))\n",
        "  samples_99_percentile = np.percentile(np.abs(data), 99.9)\n",
        "  normalized_samples = data / samples_99_percentile\n",
        "  normalized_samples = np.clip(normalized_samples, -1, 1)\n",
        "  return normalized_samples\n",
        "\n",
        "class EmbeddingDataFileList(object):\n",
        "  \"\"\"Container that loads audio, stores it as embeddings and can\n",
        "  rebalance it.\"\"\"\n",
        "\n",
        "  def __init__(self, filelist,\n",
        "               data_dest_dir,\n",
        "               targets=None,\n",
        "               label_max=10000,\n",
        "               negative_label=\"negative\",\n",
        "               negative_multiplier=25,\n",
        "               target_samples=32000,\n",
        "               progress_bar=None,\n",
        "               embedding_model=None):\n",
        "    \"\"\"Creates an instance of `EmbeddingDataFileList`.\"\"\"\n",
        "    self._negative_label = negative_label\n",
        "    self._data_per_label = collections.defaultdict(list)\n",
        "    self._labelcounts = {}\n",
        "    self._label_list = targets\n",
        "    total_examples = sum([min(len(x), label_max) for x in filelist.values()])\n",
        "    total_examples -= min(len(filelist[negative_label]), label_max)\n",
        "    total_examples += min(len(filelist[negative_label]), negative_multiplier * label_max)\n",
        "    print(\"loading %d examples\" % total_examples)\n",
        "    example_count = 0\n",
        "    for label in filelist:\n",
        "      if label not in self._label_list:\n",
        "        raise ValueError(\"Unknown label:\", label)\n",
        "      label_files = filelist[label]\n",
        "      random.shuffle(label_files)\n",
        "      if label == negative_label:\n",
        "        multplier = negative_multiplier\n",
        "      else:\n",
        "        multplier = 1\n",
        "      for wav_file in label_files[:label_max * multplier]:\n",
        "        data = normalized_read(os.path.join(data_dest_dir, wav_file))\n",
        "        required_padding = target_samples - data.shape[0]\n",
        "        if required_padding \u003e 0:\n",
        "          data = np.pad(data, (required_padding, required_padding), 'constant')\n",
        "        self._labelcounts[label] = self._labelcounts.get(label, 0) + 1\n",
        "        if embedding_model:\n",
        "           data = embedding_model.create_embedding(data)[0][0,:,:,:]\n",
        "        self._data_per_label[label].append(data)\n",
        "        if progress_bar is not None:\n",
        "          example_count += 1\n",
        "          progress_bar.update(progress(100 * example_count/total_examples))\n",
        "\n",
        "  @property\n",
        "  def labels(self):\n",
        "    return self._label_list\n",
        "\n",
        "  def get_label(self, idx):\n",
        "    return self.labels.index(idx)\n",
        "\n",
        "  def _get_filtered_data(self, label, filter_fn):\n",
        "    idx = self.labels.index(label)\n",
        "    return [(filter_fn(x), idx) for x in self._data_per_label[label]]\n",
        "\n",
        "  def _multply_data(self, data, factor):\n",
        "    samples = int((factor - math.floor(factor)) * len(data))\n",
        "    return int(factor) * data + random.sample(data, samples)\n",
        "\n",
        "  def full_rebalance(self, negatives, labeled):\n",
        "    \"\"\"Rebalances for a given ratio of labeled to negatives.\"\"\"\n",
        "    negative_count = self._labelcounts[self._negative_label]\n",
        "    labeled_count = sum(self._labelcounts[key]\n",
        "                        for key in self._labelcounts.keys()\n",
        "                        if key != self._negative_label)\n",
        "    labeled_multiply = labeled * negative_count / (negatives * labeled_count)\n",
        "    for label in self._data_per_label:\n",
        "      if label == self._negative_label:\n",
        "        continue\n",
        "      self._data_per_label[label] = self._multply_data(\n",
        "          self._data_per_label[label], labeled_multiply)\n",
        "      self._labelcounts[label] = len(self._data_per_label[label])\n",
        "\n",
        "  def get_all_data_shuffled(self, filter_fn):\n",
        "    \"\"\"Returns a shuffled list containing all the data.\"\"\"\n",
        "    return self.get_all_data(filter_fn, shuffled=True)\n",
        "\n",
        "  def get_all_data(self, filter_fn, shuffled=False):\n",
        "    \"\"\"Returns a list containing all the data.\"\"\"\n",
        "    data = []\n",
        "    for label in self._data_per_label:\n",
        "      data += self._get_filtered_data(label, filter_fn)\n",
        "    if shuffled:\n",
        "      random.shuffle(data)\n",
        "    return data\n",
        "\n",
        "def cut_middle_frame(embedding, num_frames, flatten):\n",
        "  \"\"\"Extrats the middle frames for an embedding.\"\"\"\n",
        "  left_context = (embedding.shape[0] - num_frames) // 2\n",
        "  if flatten:\n",
        "    return embedding[left_context:left_context+num_frames].flatten()\n",
        "  else:\n",
        "    return embedding[left_context:left_context+num_frames]\n",
        "\n",
        "\n",
        "def progress(value, maximum=100):\n",
        "  return IPython.display.HTML(\"\"\"\n",
        "  \u003cprogress value='{value}' max='{max}' style='width: 80%'\u003e{value}\u003c/progress\u003e\n",
        "    \"\"\".format(value=value, max=maximum))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "rEuL8h_GztBj"
      },
      "outputs": [],
      "source": [
        "#@title HeadTrainerClass and head model functions\n",
        "\n",
        "def _fully_connected_model_fn(embeddings, num_labels):\n",
        "  \"\"\"Builds the head model and adds a fully connected output layer.\"\"\"\n",
        "  net = tf.layers.flatten(embeddings)\n",
        "  logits = tf.compat.v1.layers.dense(net, num_labels, activation=None)\n",
        "  return logits\n",
        "\n",
        "framework = tf.contrib.framework\n",
        "layers = tf.contrib.layers\n",
        "\n",
        "def _conv_head_model_fn(embeddings, num_labels, context):\n",
        "  \"\"\"Builds the head model and adds a fully connected output layer.\"\"\"\n",
        "  activation_fn = tf.nn.elu\n",
        "  normalizer_fn = functools.partial(\n",
        "      layers.batch_norm, scale=True, is_training=True)\n",
        "  with framework.arg_scope([layers.conv2d], biases_initializer=None,\n",
        "                           activation_fn=None, stride=1, padding=\"SAME\"):\n",
        "    net = embeddings\n",
        "    net = layers.conv2d(net, 96, [3, 1])\n",
        "    net = normalizer_fn(net)\n",
        "    net = activation_fn(net)\n",
        "    net = layers.max_pool2d(net, [2, 1], stride=[2, 1], padding=\"VALID\")\n",
        "    context //= 2\n",
        "    net = layers.conv2d(net, 96, [3, 1])\n",
        "    net = normalizer_fn(net)\n",
        "    net = activation_fn(net)\n",
        "    net = layers.max_pool2d(net, [context, net.shape[2]], padding=\"VALID\")\n",
        "  net = tf.layers.flatten(net)\n",
        "  logits = layers.fully_connected(\n",
        "      net, num_labels, activation_fn=None)\n",
        "  return logits\n",
        "\n",
        "class HeadTrainer(object):\n",
        "  \"\"\"A tensorflow classifier to quickly train and test on embeddings.\n",
        "\n",
        "  Only use this if you are training a very small model on a very limited amount\n",
        "  of data. If you expect the training to take any more than 15 - 20 min then use\n",
        "  something else.\n",
        "  \"\"\"\n",
        "\n",
        "  def __init__(self, model_fn, input_shape, num_targets,\n",
        "               head_learning_rate=0.001, batch_size=64):\n",
        "    \"\"\"Creates a `HeadTrainer`.\n",
        "\n",
        "    Args:\n",
        "      model_fn: function that builds the tensorflow model, defines its loss\n",
        "          and returns the tuple (predictions, loss, accuracy).\n",
        "      input_shape: describes the shape of the models input feature.\n",
        "          Does not include a the batch dimension.\n",
        "      num_targets: Target number of keywords.\n",
        "    \"\"\"\n",
        "    self._input_shape = input_shape\n",
        "    self._output_dim = num_targets\n",
        "    self._batch_size = batch_size\n",
        "    self._graph = tf.Graph()\n",
        "    with self._graph.as_default():\n",
        "      self._feature = tf.placeholder(tf.float32, shape=([None] + input_shape))\n",
        "      self._labels = tf.placeholder(tf.int64, shape=(None))\n",
        "      module_spec = hub.create_module_spec(\n",
        "          module_fn=self._get_headmodule_fn(model_fn, num_targets))\n",
        "      self._module = hub.Module(module_spec, trainable=True)\n",
        "      logits = self._module(self._feature)\n",
        "      self._predictions = tf.nn.softmax(logits)\n",
        "      self._loss, self._accuracy = self._get_loss(\n",
        "          logits, self._labels, self._predictions)\n",
        "      self._update_weights = tf.train.AdamOptimizer(\n",
        "          learning_rate=head_learning_rate).minimize(self._loss)\n",
        "    self._sess = tf.Session(graph=self._graph)\n",
        "    with self._sess.as_default():\n",
        "      with self._graph.as_default():\n",
        "        self._sess.run(tf.local_variables_initializer())\n",
        "        self._sess.run(tf.global_variables_initializer())\n",
        "\n",
        "  def _get_headmodule_fn(self, model_fn, num_targets):\n",
        "    \"\"\"Wraps the model_fn in a tf hub module.\"\"\"\n",
        "    def module_fn():\n",
        "      embeddings = tf.placeholder(\n",
        "          tf.float32, shape=([None] + self._input_shape))\n",
        "      logit = model_fn(embeddings, num_targets)\n",
        "      hub.add_signature(name='default', inputs=embeddings, outputs=logit)\n",
        "    return module_fn\n",
        "\n",
        "\n",
        "  def _get_loss(self, logits, labels, predictions):\n",
        "    \"\"\"Defines the model's loss and accuracy.\"\"\"\n",
        "    xentropy_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(\n",
        "        logits=logits, labels=labels)\n",
        "    loss = tf.reduce_mean(xentropy_loss)\n",
        "    accuracy = tf.contrib.metrics.accuracy(tf.argmax(predictions, 1), labels)\n",
        "    return loss, accuracy\n",
        "\n",
        "  def save_head_model(self, save_directory):\n",
        "    \"\"\"Saves the model.\"\"\"\n",
        "    with self._graph.as_default():\n",
        "      self._module.export(save_directory, self._sess)\n",
        "\n",
        "\n",
        "  def _feature_transform(self, batch_features, batch_labels):\n",
        "    \"\"\"Transforms lists of features and labels into into model inputs.\"\"\"\n",
        "    return np.stack(batch_features), np.stack(batch_labels)\n",
        "\n",
        "  def _batch_data(self, data, batch_size=None):\n",
        "    \"\"\"Splits the input data into batches.\"\"\"\n",
        "    batch_features = []\n",
        "    batch_labels = []\n",
        "    batch_size = batch_size or len(data)\n",
        "    for feature, label in data:\n",
        "      if feature.shape != tuple(self._input_shape):\n",
        "        raise ValueError(\n",
        "            \"Feature shape ({}) doesn't match model shape ({})\".format(\n",
        "                feature.shape, self._input_shape))\n",
        "      if not 0 \u003c= label \u003c self._output_dim:\n",
        "\n",
        "        raise ValueError('Label value ({}) outside of target range'.format(\n",
        "            label))\n",
        "      batch_features.append(feature)\n",
        "      batch_labels.append(label)\n",
        "      if len(batch_features) == batch_size:\n",
        "        yield self._feature_transform(batch_features, batch_labels)\n",
        "        del batch_features[:]\n",
        "        del batch_labels[:]\n",
        "    if batch_features:\n",
        "      yield self._feature_transform(batch_features, batch_labels)\n",
        "\n",
        "  def epoch_train(self, data, epochs=1, batch_size=None):\n",
        "    \"\"\"Trains the model on the provided data.\n",
        "\n",
        "    Args:\n",
        "      data: List of tuples (feature, label) where feature is a np array of\n",
        "          shape `self._input_shape` and label an int less than self._output_dim.\n",
        "      epochs: Number of times this data should be trained on.\n",
        "      batch_size: Number of feature, label pairs per batch. Overwrites\n",
        "          `self._batch_size` when set.\n",
        "\n",
        "    Returns:\n",
        "      tuple of accuracy, loss;\n",
        "          accuracy: Average training accuracy.\n",
        "          loss: Loss of the final batch.\n",
        "    \"\"\"\n",
        "    batch_size = batch_size or self._batch_size\n",
        "    accuracy_list = []\n",
        "    for _ in range(epochs):\n",
        "      for features, labels in self._batch_data(data, batch_size):\n",
        "        loss, accuracy, _ = self._sess.run(\n",
        "            [self._loss, self._accuracy, self._update_weights],\n",
        "            feed_dict={self._feature: features, self._labels: labels})\n",
        "        accuracy_list.append(accuracy)\n",
        "    return (sum(accuracy_list))/len(accuracy_list), loss\n",
        "\n",
        "  def test(self, data, batch_size=None):\n",
        "    \"\"\"Evaluates the model on the provided data.\n",
        "\n",
        "    Args:\n",
        "      data: List of tuples (feature, label) where feature is a np array of\n",
        "          shape `self._input_shape` and label an int less than self._output_dim.\n",
        "      batch_size: Number of feature, label pairs per batch. Overwrites\n",
        "          `self._batch_size` when set.\n",
        "\n",
        "    Returns:\n",
        "      tuple of accuracy, loss;\n",
        "          accuracy: Average training accuracy.\n",
        "          loss: Loss of the final batch.\n",
        "    \"\"\"\n",
        "    batch_size = batch_size or self._batch_size\n",
        "    accuracy_list = []\n",
        "    for features, labels in self._batch_data(data, batch_size):\n",
        "      loss, accuracy = self._sess.run(\n",
        "          [self._loss, self._accuracy],\n",
        "          feed_dict={self._feature: features, self._labels: labels})\n",
        "      accuracy_list.append(accuracy)\n",
        "    return sum(accuracy_list)/len(accuracy_list), loss\n",
        "\n",
        "  def infer(self, example_feature):\n",
        "    \"\"\"Runs inference on example_feature.\"\"\"\n",
        "    if example_feature.shape != tuple(self._input_shape):\n",
        "      raise ValueError(\n",
        "          \"Feature shape ({}) doesn't match model shape ({})\".format(\n",
        "              example_feature.shape, self._input_shape))\n",
        "    return self._sess.run(\n",
        "        self._predictions,\n",
        "        feed_dict={self._feature: np.expand_dims(example_feature, axis=0)})"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "FVLlx0xSztBn"
      },
      "outputs": [],
      "source": [
        "#@title TfHubWrapper Class\n",
        "\n",
        "class TfHubWrapper(object):\n",
        "  \"\"\"A loads a tf hub embedding model.\"\"\"\n",
        "  def __init__(self, embedding_model_dir):\n",
        "    \"\"\"Creates a `SavedModelWraper`.\"\"\"\n",
        "    self._graph = tf.Graph()\n",
        "    self._sess = tf.Session(graph=self._graph)\n",
        "    with self._graph.as_default():\n",
        "      with self._sess.as_default():\n",
        "        module_spec = hub.load_module_spec(embedding_model_dir)\n",
        "        embedding_module = hub.Module(module_spec)\n",
        "        self._samples = tf.placeholder(\n",
        "            tf.float32, shape=[1, None], name='audio_samples')\n",
        "        self._embedding = embedding_module(self._samples)\n",
        "        self._sess.run(tf.global_variables_initializer())\n",
        "    print(\"Embedding model loaded, embedding shape:\", self._embedding.shape)\n",
        "\n",
        "  def create_embedding(self, samples):\n",
        "    samples = samples.reshape((1, -1))\n",
        "    output = self._sess.run(\n",
        "        [self._embedding],\n",
        "        feed_dict={self._samples: samples})\n",
        "    return output"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "eMlokdbTZ4_n"
      },
      "outputs": [],
      "source": [
        "#@title Define AudioClipRecorder Class\n",
        "AUDIOCLIP_HTML ='''\n",
        "\u003cspan style=\"font-size:30px\"\u003eRecorded audio clips of {keyphrase}:\u003c/span\u003e \n",
        "\u003cdiv id='target{keyphrase}'\u003e\u003c/div\u003e\n",
        "\u003cspan id = \"status_label{keyphrase}\" style=\"font-size:30px\"\u003e\n",
        "  Ready to record.\u003c/span\u003e\n",
        "\u003cbutton id='Add{keyphrase}Audio'\u003eRecord\u003c/button\u003e\n",
        "\u003cscript\u003e\n",
        "var recorder;\n",
        "var base64data = 0;\n",
        "\n",
        "function sleep(ms) {{\n",
        "  return new Promise(resolve =\u003e setTimeout(resolve, ms));\n",
        "}}\n",
        "\n",
        "var handleSuccess = function(stream) {{\n",
        "  recorder = new MediaRecorder(stream);\n",
        "  recorder.ondataavailable = function(e) {{            \n",
        "    reader = new FileReader();\n",
        "    reader.readAsDataURL(e.data); \n",
        "    reader.onloadend = function() {{\n",
        "      base64data = reader.result;\n",
        "    }}\n",
        "  }};\n",
        "  recorder.start();\n",
        "}};\n",
        "\n",
        "document.querySelector('#Add{keyphrase}Audio').onclick = () =\u003e {{\n",
        "  var label = document.getElementById(\"status_label{keyphrase}\"); \n",
        "  navigator.mediaDevices.getUserMedia({{audio: true}}).then(handleSuccess);\n",
        "  label.innerHTML = \"Recording ... please say {keyphrase}!\".fontcolor(\"red\");; \n",
        "    sleep({clip_length_ms}).then(() =\u003e {{\n",
        "    recorder.stop();\n",
        "    label.innerHTML = \"Recording finished ... processing audio.\"; \n",
        "    sleep(1000).then(() =\u003e {{\n",
        "      google.colab.kernel.invokeFunction('notebook.AddAudioItem{keyphrase}',\n",
        "      [base64data.toString()], {{}});\n",
        "      label.innerHTML = \"Ready to record.\";\n",
        "    }});\n",
        "}});\n",
        "}};\n",
        "\u003c/script\u003e'''\n",
        "\n",
        "class AudioClipRecorder:\n",
        "  \"\"\"Python class that creates a JS microphone clip recorder.\"\"\"\n",
        "\n",
        "  def __init__(self, keyphrase=\"test\", clip_length_ms=2100):\n",
        "    \"\"\"Creates an AudioClipRecorder instance.\n",
        "\n",
        "    When created this class prints an empty \u003cdiv\u003e tag into which the\n",
        "    recorded clips will be printed and a record audio button that uses\n",
        "    javascript to access the microphone and record an audio clip.\n",
        "    \n",
        "    Args:\n",
        "      keyphrase: The name of the keyphrase that should be recorded.\n",
        "        This will be displayed in the recording prompt and used as a\n",
        "        directory name when the recordings are exported.\n",
        "      clip_length_ms: The length (in ms) of each recorded audio clip.\n",
        "        Due to the async nature of javascript this actual amount of recorded\n",
        "        audio may vary by a ~20-80ms.\n",
        "    \"\"\"\n",
        "    self._counter = 0\n",
        "    self._keyphrase = keyphrase\n",
        "    self._audio_clips = {}\n",
        "    IPython.display.display(IPython.display.HTML(AUDIOCLIP_HTML.format(\n",
        "        keyphrase=keyphrase, clip_length_ms=clip_length_ms)))\n",
        "    output.register_callback('notebook.AddAudioItem' + keyphrase,\n",
        "                             self.add_list_item)\n",
        "    output.register_callback('notebook.RemoveAudioItem' + keyphrase,\n",
        "                             self.rm_audio)\n",
        "\n",
        "  def add_list_item(self, data):\n",
        "    \"\"\"Adds the recorded audio to the list of clips.\n",
        "\n",
        "    This function is called from javascript after clip_length_ms audio has\n",
        "    been recorded. It prints the recorded audio clip to the \u003cdiv\u003e together with\n",
        "    a button that allows for it to be deleted.\n",
        "\n",
        "    Args:\n",
        "      data: The recorded audio in webm format.\n",
        "    \"\"\"\n",
        "    raw_string_data = data.split(',')[1]\n",
        "    samples, rate = self.decode_webm(raw_string_data)\n",
        "    length_samples = len(samples)\n",
        "    with output.redirect_to_element('#target{keyphrase}'.format(\n",
        "        keyphrase=self._keyphrase)):\n",
        "      with output.use_tags('{keyphrase}_audio_{counter}'.format(\n",
        "          counter=self._counter, keyphrase=self._keyphrase)):\n",
        "        IPython.display.display(IPython.display.HTML('''Audio clip {counter} - \n",
        "        {length} samples - \n",
        "        \u003cbutton id=\\'delbutton{keyphrase}{counter}\\'\u003edel\u003c/button\u003e\n",
        "        \u003cscript\u003e\n",
        "        document.querySelector('#delbutton{keyphrase}{counter}').onclick = () =\u003e {{\n",
        "          google.colab.kernel.invokeFunction('notebook.RemoveAudioItem{keyphrase}', [{counter}], {{}});\n",
        "        }};\n",
        "        \u003c/script\u003e'''.format(counter=self._counter, length=length_samples,\n",
        "                            keyphrase=self._keyphrase)))\n",
        "        IPython.display.display(IPython.display.Audio(data=samples, rate=rate))\n",
        "        IPython.display.display(IPython.display.HTML('\u003cbr\u003e\u003cbr\u003e'))\n",
        "        self._audio_clips[self._counter]=samples\n",
        "      self._counter+=1\n",
        "\n",
        "  def rm_audio(self, count):\n",
        "    \"\"\"Removes the audioclip 'count' from the list of clips.\"\"\"\n",
        "    output.clear(output_tags=\"{0}_audio_{1}\".format(self._keyphrase, count))\n",
        "    self._audio_clips.pop(count)\n",
        "\n",
        "  def decode_webm(self, data):\n",
        "    \"\"\"Decodes a webm audio clip in a np.array of samples.\"\"\"\n",
        "    sample_rate=16000\n",
        "    process = (ffmpeg\n",
        "      .input('pipe:0')\n",
        "      .output('pipe:1', format='s16le', ar=sample_rate)\n",
        "      .run_async(pipe_stdin=True, pipe_stdout=True, pipe_stderr=True,\n",
        "                 quiet=True, overwrite_output=True)\n",
        "    )\n",
        "    output, err = process.communicate(input=b64decode(data))\n",
        "    audio = np.frombuffer(output, dtype=np.int16).astype(np.float32)\n",
        "    return audio, sample_rate\n",
        "\n",
        "  def save_as_wav_files(self, base_output_dir,\n",
        "                        file_prefix='recording_', file_suffix=''):\n",
        "    \"\"\"Exports all audio clips as wav files.\n",
        "\n",
        "    The files wav files will be written to 'base_output_dir/self._keyphrase'.\n",
        "    And will be named: file_prefix + str(clip_id) + file_suffix + '.wav'    \n",
        "    \"\"\"\n",
        "    if not os.path.exists(base_output_dir):\n",
        "      os.mkdir(base_output_dir)\n",
        "    keyphrase_output_dir = os.path.join(base_output_dir, self._keyphrase)\n",
        "    if not os.path.exists(keyphrase_output_dir):\n",
        "      os.mkdir(keyphrase_output_dir)\n",
        "    for clip_id in self._audio_clips:\n",
        "      filename = file_prefix + str(clip_id) + file_suffix + '.wav'\n",
        "      output_file = os.path.join(keyphrase_output_dir, filename)\n",
        "      print(\"Creating:\", output_file)\n",
        "      scipy.io.wavfile.write(output_file, 16000, self._audio_clips[clip_id])\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "GdyxtK-mrO8l"
      },
      "outputs": [],
      "source": [
        "#@title Define AudioClipEval Class\n",
        "class AudioClipEval(AudioClipRecorder):\n",
        "  def __init__(self, embedding_model, head_model, filter_fn, labels,\n",
        "               name=\"eval1\", clip_length_ms=2100):\n",
        "    \"\"\"Creates an AudioClipEval instance.\n",
        "\n",
        "    When created this class prints an empty \u003cdiv\u003e tag into which the\n",
        "    recorded clips will be printed and a record audio button that uses\n",
        "    javascript to access the microphone and record an audio clip.\n",
        "    \n",
        "    Args:\n",
        "      embedding_model: The embedding model.\n",
        "      head_model: The default head model.\n",
        "      filter_fn: function that prepared the input to the head model.\n",
        "      labels: List of head model target labels.\n",
        "      keyphrase: The name of the keyphrase that should be recorded.\n",
        "        This will be displayed in the recording prompt and used as a\n",
        "        directory name when the recordings are exported.\n",
        "      clip_length_ms: The length (in ms) of each recorded audio clip.\n",
        "        Due to the async nature of javascript this actual amount of recorded\n",
        "        audio may vary by a ~20-80ms.\n",
        "    \"\"\"\n",
        "    self._counter = 0\n",
        "    self._keyphrase = name\n",
        "    keyphrase = name\n",
        "    self._audio_clips = {}\n",
        "    self._embedding_model = embedding_model\n",
        "    self._head_model = head_model\n",
        "    self._filter_fn = filter_fn\n",
        "    self._labels = labels\n",
        "    IPython.display.display(IPython.display.HTML(\n",
        "        AUDIOCLIP_HTML.format(keyphrase=keyphrase, clip_length_ms=clip_length_ms)))\n",
        "    output.register_callback('notebook.AddAudioItem' + keyphrase,\n",
        "                             self.add_list_item)\n",
        "    output.register_callback('notebook.RemoveAudioItem' + keyphrase,\n",
        "                             self.rm_audio)\n",
        "\n",
        "  def add_list_item(self, data):\n",
        "    \"\"\"Adds the recorded audio to the list of clips and classifies it.\n",
        "\n",
        "    This function is called from javascript after clip_length_ms audio has\n",
        "    been recorded. It prints the recorded audio clip to the \u003cdiv\u003e together with\n",
        "    a button that allows for it to be deleted.\n",
        "\n",
        "    Args:\n",
        "      data: The recorded audio in webm format.\n",
        "    \"\"\"\n",
        "    raw_string_data = data.split(',')[1]\n",
        "    samples, rate = self.decode_webm(raw_string_data)\n",
        "    length_samples = len(samples)\n",
        "    detection, confidence = self.eval_audio(samples)\n",
        "    with output.redirect_to_element('#target{keyphrase}'.format(\n",
        "        keyphrase=self._keyphrase)):\n",
        "      with output.use_tags('{keyphrase}_audio_{counter}'.format(\n",
        "          counter=self._counter, keyphrase=self._keyphrase)):\n",
        "        IPython.display.display(IPython.display.HTML('''Audio clip {counter} - \n",
        "        {length} samples - \n",
        "        \u003cbutton id=\\'delbutton{counter}\\'\u003edel\u003c/button\u003e\n",
        "        \u003cscript\u003e\n",
        "        document.querySelector('#delbutton{counter}').onclick = () =\u003e {{\n",
        "          google.colab.kernel.invokeFunction('notebook.RemoveAudioItem{keyphrase}', [{counter}], {{}});\n",
        "        }};\n",
        "        \u003c/script\u003e'''.format(counter=self._counter, length=length_samples,\n",
        "                            keyphrase=self._keyphrase)))\n",
        "        IPython.display.display(IPython.display.Audio(data=samples, rate=rate))\n",
        "        IPython.display.display(IPython.display.HTML(\n",
        "            '''\u003cspan id = \"result{counter}\" style=\"font-size:24px\"\u003e\n",
        "              detected: {detection} ({confidence})\u003cspan\u003e'''.format(\n",
        "                  counter=self._counter, detection=detection,\n",
        "                  confidence=confidence)))\n",
        "        IPython.display.display(IPython.display.HTML('\u003cbr\u003e\u003cbr\u003e'))\n",
        "        self._audio_clips[self._counter]=samples\n",
        "      self._counter+=1\n",
        "\n",
        "  def eval_audio(self, samples, head_model=None):\n",
        "    \"\"\"Classifies the audio using the current or a provided model.\"\"\"\n",
        "    embeddings = self._embedding_model.create_embedding(samples)[0][0,:,:,:]\n",
        "    if head_model:\n",
        "      probs = head_model.infer(self._filter_fn(embeddings))\n",
        "    else:\n",
        "      probs = self._head_model.infer(self._filter_fn(embeddings))\n",
        "    return self._labels[np.argmax(probs)], np.amax(probs)\n",
        "\n",
        "  def eval_on_new_model(self, head_model):\n",
        "    \"\"\"Reclassifies the clips using a new head model.\"\"\"\n",
        "    for clip_id in self._audio_clips:\n",
        "      samples = self._audio_clips[clip_id]\n",
        "      length_samples = len(samples)\n",
        "      detection, confidence = self.eval_audio(samples, head_model=head_model)\n",
        "      IPython.display.display(IPython.display.HTML(\n",
        "            '''Audio clip {counter} -  {length} samples - \n",
        "            \u003cspan id = \"result{counter}\" style=\"font-size:24px\"\u003e\n",
        "              detected: {detection} ({confidence})\u003cspan\u003e'''.format(\n",
        "                counter=clip_id, length=length_samples,\n",
        "                detection=detection, confidence=confidence)))     \n",
        "      IPython.display.display(IPython.display.Audio(data=samples, rate=16000))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "gR6n9PMGIHSv"
      },
      "source": [
        "## Load the embedding model\n",
        "\n",
        "The following info messages can be ignored\n",
        "\n",
        "\u003e *INFO:tensorflow:Saver not created because there are no variables in the graph to restore*\n",
        "\n",
        "Don't worry tf hub is restoring all the variables.\n",
        "\n",
        "You can test the model by having it produce an embedding on zeros:\n",
        "\n",
        "\n",
        "```\n",
        "speech_embedding_model.create_embedding(np.zeros((1,66000)))\n",
        "```\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "O4LRwMdsIVpo"
      },
      "outputs": [],
      "source": [
        "embedding_model_url = \"https://tfhub.dev/google/speech_embedding/1\"\n",
        "speech_embedding_model = TfHubWrapper(embedding_model_url)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "omDLyJkac0RH"
      },
      "source": [
        "## Record training data or copy from google drive\n",
        "\n",
        "The following cells allow you to define a set of target keyphrases and record some examples for training.\n",
        "\n",
        "### Optional Google Drive access.\n",
        "\n",
        "The recorded wav files can be uploaded (and later download) from your Google drive using [PyDrive](https://googleworkspace.github.io/PyDrive/docs/build/html/index.html). When you run the *Set up Google drive access* cell it will prompt you to log in and grant this colab permission to access your Google drive. Only if you do this will you be able to run the other Google drive cells.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "sdTXzHMSO5U5"
      },
      "outputs": [],
      "source": [
        "#@title Optional: Set up Google drive access\n",
        "!pip install PyDrive\n",
        "from pydrive.auth import GoogleAuth\n",
        "from pydrive.drive import GoogleDrive\n",
        "from google.colab import auth\n",
        "from oauth2client.client import GoogleCredentials\n",
        "auth.authenticate_user()\n",
        "gauth = GoogleAuth()\n",
        "gauth.credentials = GoogleCredentials.get_application_default()\n",
        "drive = GoogleDrive(gauth)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "G7cs8_fiN4-W"
      },
      "outputs": [],
      "source": [
        "#@title Optional: Download and untar an archive from drive\n",
        "\n",
        "filename = ''#@param {type:\"string\"}\n",
        "#@markdown You can find the file_id by looking at its share-link.\n",
        "#@markdown e.g. *1b9Lkfie2NHX-O06vPGrqzyGcGWUPul36*\n",
        "file_id =  ''#@param {type:\"string\"}\n",
        "\n",
        "downloaded = drive.CreateFile({'id':file_id})\n",
        "downloaded.GetContentFile(filename)\n",
        "with tarfile.open(filename, 'r:gz') as data_tar_file:\n",
        "  for member_info in data_tar_file.getmembers():\n",
        "    print(member_info.name)\n",
        "    data_tar_file.extract(member_info)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "zKdnYUSpcY0G"
      },
      "outputs": [],
      "source": [
        "#@title Setup recording session and define model targets\n",
        "\n",
        "#@markdown Only use letters and _ for the **RECORDING_NAME** and **TARGET_WORDS**. \n",
        "RECORDING_NAME = 'transportation' #@param {type:\"string\"}\n",
        "target_word1 = 'hogwarts_express'  #@param {type:\"string\"}\n",
        "target_word2 = 'nautilus'  #@param {type:\"string\"}\n",
        "target_word3 = 'millennium_falcon'  #@param {type:\"string\"}\n",
        "target_word4 = 'enterprise'  #@param {type:\"string\"}\n",
        "target_word5 = ''  #@param {type:\"string\"}\n",
        "target_word6 = ''  #@param {type:\"string\"}\n",
        "clip_lengh_ms = 2100 #@param {type:\"integer\"}\n",
        "\n",
        "#@markdown ### Microphone access\n",
        "#@markdown Please connect the microphone that you want to use\n",
        "#@markdown before running this cell. You may also be asked to\n",
        "#@markdown to grant colab permission to use it.\n",
        "#@markdown If you have any problems check your browser settings\n",
        "#@markdown and rerun the cell.\n",
        "\n",
        "target_words = [target_word1, target_word2, target_word3,\n",
        "                target_word4, target_word5, target_word6]\n",
        "\n",
        "OWN_TARGET_WORDS = ','.join([w for w in target_words if w is not ''])\n",
        "OWN_MODEL_LABELS = ['negative', 'silence'] + OWN_TARGET_WORDS.split(',')\n",
        "\n",
        "word_list = OWN_TARGET_WORDS.split(',')\n",
        "\n",
        "t = widgets.TabBar(word_list)\n",
        "\n",
        "clip_recorders = {}\n",
        "for label in word_list:\n",
        "  with t.output_to(word_list.index(label)):\n",
        "    clip_recorders[label] = AudioClipRecorder(keyphrase=label,\n",
        "                                              clip_length_ms=2100)\n",
        "\n",
        "with t.output_to(0):\n",
        "  print()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "xLhKuO9eUljA"
      },
      "outputs": [],
      "source": [
        "#@title Create wav files from recording session.\n",
        "\n",
        "session =  'recording1_'#@param {type:\"string\"}\n",
        "speaker =  '_spk1'#@param {type:\"string\"}\n",
        "\n",
        "for label in clip_recorders:\n",
        "  clip_recorders[label].save_as_wav_files(base_output_dir=RECORDING_NAME,\n",
        "                                          file_prefix=session,\n",
        "                                          file_suffix=speaker)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "3RJnwChkgl_6"
      },
      "outputs": [],
      "source": [
        "#@title Load files for training.\n",
        "                                    \n",
        "all_train_example_files = collections.defaultdict(list)\n",
        "\n",
        "for label in OWN_TARGET_WORDS.split(','):\n",
        "  label_dir = os.path.join(RECORDING_NAME, label)\n",
        "  all_label_files = [\n",
        "      os.path.join(label, f)\n",
        "      for f in os.listdir(label_dir)\n",
        "      if os.path.isfile(os.path.join(label_dir, f))\n",
        "  ]\n",
        "  all_train_example_files[label].extend(all_label_files)\n",
        "\n",
        "progress_bar = IPython.display.display(progress(0, 100), display_id=True)\n",
        "print(\"loading train data\")\n",
        "train_data = EmbeddingDataFileList(\n",
        "    all_train_example_files, RECORDING_NAME,\n",
        "    targets=OWN_MODEL_LABELS, embedding_model=speech_embedding_model,\n",
        "    progress_bar=progress_bar)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "FtWvEIPDGKhq"
      },
      "outputs": [],
      "source": [
        "#@title Optional: save recorded data to drive.\n",
        "\n",
        "archive_name = RECORDING_NAME + \"_\" + str(int(time.time())) +\".tar.gz\"\n",
        "\n",
        "def make_tarfile(output_filename, source_dir):\n",
        "    with tarfile.open(output_filename, \"w:gz\") as tar:\n",
        "        tar.add(source_dir, arcname=os.path.basename(source_dir))\n",
        "\n",
        "make_tarfile(archive_name, RECORDING_NAME)\n",
        "\n",
        "file1 = drive.CreateFile({'title': archive_name})\n",
        "file1.SetContentFile(archive_name)\n",
        "file1.Upload()\n",
        "print('Saving to drive: %s, id: %s' % (file1['title'], file1['id']))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "A1bAhd8U4h70"
      },
      "source": [
        "# Train a model on your recorded data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "79le-IpzVgsr"
      },
      "outputs": [],
      "source": [
        "#@title Run training\n",
        "\n",
        "#@markdown We assume that the keyphrase is spoken roughly in the middle\n",
        "#@markdown of the loaded audio clips. With **context_size** we can choose the \n",
        "#@markdown number of embeddings around the middle to use as a model input.\n",
        "context_size = 16 #@param {type:\"slider\", min:1, max:28, step:1}\n",
        "\n",
        "filter_fn = functools.partial(cut_middle_frame, num_frames=context_size, flatten=False)\n",
        "all_train_data = train_data.get_all_data_shuffled(filter_fn=filter_fn)\n",
        "all_eval_data = None\n",
        "\n",
        "head_model = \"Convolutional\" #@param [\"Convolutional\", \"Fully_Connected\"] {type:\"string\"}\n",
        "\n",
        "#@markdown Suggested **learning_rate** range: 0.00001 - 0.01.\n",
        "learning_rate = 0.001 #@param {type:\"number\"}\n",
        "batch_size = 32\n",
        "#@markdown **epochs_per_eval** and **train_eval_loops** control how long the\n",
        "#@markdown the model is trained. An epoch is defined as the model having seen\n",
        "#@markdown each example at least once, with some examples twice to ensure the\n",
        "#@markdown correct labeled / negatives balance.\n",
        "\n",
        "epochs_per_eval = 1 #@param {type:\"slider\", min:1, max:15, step:1}\n",
        "train_eval_loops = 30 #@param {type:\"slider\", min:5, max:80, step:5}\n",
        "\n",
        "if head_model == \"Convolutional\":\n",
        "  model_fn = functools.partial(_conv_head_model_fn, context=context_size)\n",
        "else:\n",
        "  model_fn = _fully_connected_model_fn\n",
        "\n",
        "trainer = HeadTrainer(model_fn=model_fn,\n",
        "                      input_shape=[context_size,1,96],\n",
        "                      num_targets=len(OWN_MODEL_LABELS),\n",
        "                      head_learning_rate=learning_rate,\n",
        "                      batch_size=batch_size)\n",
        "\n",
        "data_trained_on = 0\n",
        "data = [] \n",
        "train_results = []\n",
        "eval_results = []\n",
        "max_data = len(all_train_data) * epochs_per_eval * train_eval_loops + 10\n",
        "\n",
        "def plot_step(plot, max_data, data, train_results, eval_results):\n",
        "  plot.clf()\n",
        "  plot.xlim(0, max_data)\n",
        "  plot.ylim(0.85, 1.05)\n",
        "  plot.plot(data, train_results, \"bo\")\n",
        "  plot.plot(data, train_results, \"b\", label=\"train_results\")\n",
        "  if eval_results:\n",
        "    plot.plot(data, eval_results, \"ro\")\n",
        "    plot.plot(data, eval_results, \"r\", label=\"eval_results\")\n",
        "  plot.legend(loc='lower right', fontsize=24)\n",
        "  plot.xlabel('number of examples trained on', fontsize=22)\n",
        "  plot.ylabel('Accuracy', fontsize=22)\n",
        "  plot.xticks(fontsize=20)\n",
        "  plot.yticks(fontsize=20) \n",
        "\n",
        "plt.figure(figsize=(25, 7))\n",
        "for loop in range(train_eval_loops):\n",
        "  train_accuracy, loss = trainer.epoch_train(all_train_data,\n",
        "                                             epochs=epochs_per_eval)\n",
        "  train_results.append(train_accuracy)\n",
        "  if all_eval_data:\n",
        "    eval_accuracy, loss = trainer.test(all_eval_data)\n",
        "    eval_results.append(eval_accuracy)\n",
        "  else:\n",
        "    eval_results = None\n",
        "\n",
        "  data_trained_on += len(all_train_data) * epochs_per_eval\n",
        "  data.append(data_trained_on)\n",
        "  plot_step(plt, max_data, data, train_results, eval_results)\n",
        "\n",
        "  IPython.display.display(plt.gcf())\n",
        "  if all_eval_data:\n",
        "    print(\"Highest eval accuracy: %.2f percent.\" % (100 * max(eval_results)))\n",
        "  IPython.display.clear_output(wait=True)\n",
        "\n",
        "if all_eval_data:\n",
        "  print(\"Highest eval accuracy: %.2f percent.\" % (100 * max(eval_results)))\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "3WAc7vr2sVAy"
      },
      "outputs": [],
      "source": [
        "#@title Test the model\n",
        "clip_eval = AudioClipEval(speech_embedding_model, trainer, filter_fn, OWN_MODEL_LABELS)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "TD4MhDAc0TJv"
      },
      "outputs": [],
      "source": [
        "#@title Rerun the test using a new head model (train a new head model first)\n",
        "\n",
        "clip_eval.eval_on_new_model(trainer)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "n_7IPrnuqoKR"
      },
      "source": [
        "## FAQ\n",
        "\n",
        "Q: **My model isn't very good?**\n",
        "\n",
        "A: The head model is very small and depends a lot on the initialisation weights:\n",
        " * This default setup doesn't have a negative class so it will always detect *something*. \n",
        " * Try retraining it a couple of times.\n",
        " * Reduce the learning rate a little bit.\n",
        " * Add more training examples:\n",
        "   * At 1 - 5 examples per keyphrase the model probably won't be very good.\n",
        "   * With around 10-20 examples per keyphrase it may work reasonably well; however, it may still fail to learn a keyphrase.\n",
        "   * If you only have examples from a single speaker, then it may only learn how that speaker pronounces the keyphrase.\n",
        " * Make sure your keyphrase are distinctive enough:\n",
        "   * e.g. heads up vs ketchup\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "Q: **Can I export the model and use it somewhere?**\n",
        "\n",
        "A: Yes, there's some example code in the following cells that demonstrate how that could be done. However, this simple example model is only training a between-word classifier.\n",
        "If you want to use it in any relaistic setting, you will probably also want to add:\n",
        " * A negative or non-target-word speech class: You could do this by recording 2-10 min of continuous speech that doesn't contain your target keyphrases.\n",
        " * A non-speech / silence / background-noise class: The speech commands dataset contains some examples of non-speech background audio that could be used for this, and/or you could just leave your mircophone on and record some ambient audio from the future deployement location."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "PQp-pXtr2oCG"
      },
      "source": [
        "# Export and reuse the head model\n",
        "The following cells show how the head model you just trained can be exported and reused in a graph."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "s2GxL706F-BD"
      },
      "outputs": [],
      "source": [
        "#@title Save the head model\n",
        "\n",
        "head_model_module_dir = \"head_model_module_dir\"\n",
        "trainer.save_head_model(head_model_module_dir)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "r6PfTOh7HIt1"
      },
      "outputs": [],
      "source": [
        "#@title FullModelWrapper - Example Class\n",
        "\n",
        "class FullModelWrapper(object):\n",
        "  \"\"\"A loads a save model classifier.\"\"\"\n",
        "  def __init__(self, embedding_model_dir, head_model_dir):\n",
        "    self._graph = tf.Graph()\n",
        "    self._sess = tf.Session(graph=self._graph)\n",
        "    with self._graph.as_default():\n",
        "      embedding_module_spec = hub.load_module_spec(embedding_model_dir)\n",
        "      embedding_module = hub.Module(embedding_module_spec)\n",
        "      head_module_spec = hub.load_module_spec(head_model_dir)\n",
        "      head_module = hub.Module(head_module_spec)\n",
        "      self._samples = tf.placeholder(\n",
        "          tf.float32, shape=[1, None], name='audio_samples')\n",
        "      embedding = embedding_module(self._samples)\n",
        "      logits = head_module(embedding)\n",
        "      self._predictions = tf.nn.softmax(logits)\n",
        "      with self._sess.as_default():\n",
        "        self._sess.run(tf.global_variables_initializer())\n",
        " \n",
        "  def infer(self, samples):\n",
        "    samples = samples.reshape((1, -1))\n",
        "    output = self._sess.run(\n",
        "        [self._predictions],\n",
        "        feed_dict={self._samples: samples})\n",
        "    return output"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "gS9gCV8SKIfe"
      },
      "outputs": [],
      "source": [
        "#@title Test the full model on zeros\n",
        "full_model = FullModelWrapper(embedding_model_url, head_model_module_dir)\n",
        "full_model.infer(np.zeros((1,32000)))"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "JndnmDMp66FL"
      ],
      "name": "Speech Embeddings 2: Record and train on your own data.ipynb",
      "private_outputs": true,
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
